# -*- coding: utf-8 -*-

'''
基于大模型筛选关键词
    1、下载chatGLM3模型
    2、设计prompt提问方式；
    3、判断关键词与主题词的相关性；
'''

from transformers import AutoTokenizer, AutoModel
from chatglm3.llm_chatglm import ChatGLM
import logging
logging.basicConfig(level=logging.INFO)
import re

llm = ChatGLM()

class ChatGLM3KeywordFilter:
    def __init__(self, threshold=0.5):
        """
        初始化关键词筛选器，使用 ChatGLM3 模型。
        :param model_name: 模型名称或路径
        :param threshold: 筛选关键词的相似度阈值
        """
        # self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
        # self.model = AutoModel.from_pretrained(model_name, trust_remote_code=True).half().to("cpu")  # 使用半精度加速
        self.threshold = threshold

    def clean_text(self, text):
        """清理文本，去除非中文字符"""
        text = re.sub(r'[^\u4e00-\u9fa5]', '', text)
        return text

    def find_keywords(self, theme_words, context):
        res = []
        cleaned_text = tuple(map(self.clean_text, context))
        for theme in theme_words:
            for i in range(len(cleaned_text)):
                if (i+1)%30 == 0:
                    if ''.join(cleaned_text[50*i:50*(i+1)]) != "":
                        input_text = (
                            f"根据年报上下文：{''.join(cleaned_text[50*i:50*(i+1)])}提取跟主题词：{theme}相关的关键词\n"
                            f"请输出关键词\n"
                        )
                        logging.info(f"input_text: {input_text}")

                        generated_text = llm(input_text)
                        res.append(generated_text)

        return res
        
    
    def compute_relevance(self, theme, keyword, context):
        """
        使用 ChatGLM3 计算关键词在给定上下文和主题下的相关性
        :param theme: 主题词
        :param keyword: 待筛选的关键词
        :param context: 文本上下文
        :return: 相关性分数
        """
        input_text = (
            f"主题：{theme}\n上下文：{context}\n关键词：{keyword}\n"
            f"请判断该关键词是否与主题密切相关？"
        )

        # input_text = (
        #     f"根据年报上下文：{context}提取跟主题词：{theme}相关的关键词"
        #     f"请输出关键词\n"
        #     f"请判断该关键词是否与主题密切相关？"
        # )
        logging.info(f"input_text: {input_text}")

        # self.tokenizer.tokenizer.padding_side = "left"
        # inputs = self.tokenizer(input_text, return_tensors="pt", truncation=True).to("cuda" if torch.cuda.is_available() else "cpu")
        # outputs = self.model.generate(**inputs, max_length=50)
        # generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        # model = self.model.eval()
        # generated_text = model.chat(self.tokenizer, input_text, top_p=0.95, temperature=0.5)

        generated_text = llm(input_text)
        print(generated_text)


        # 简单解析生成结果，将“是”视为高相关，“否”视为低相关
        if "是" in generated_text:
            return 1.0
        elif "否" in generated_text:
            return 0.0
        else:
            return 0.5  # 模型未明确判断时设为中等相关性

    def filter_keywords_by_theme_in_context(self, keywords, theme_words, context):
        """
        根据主题词在给定上下文中筛选关键词
        :param keywords: 待筛选的关键词列表
        :param theme_words: 主题词列表
        :param context: 文本上下文
        :return: 过滤后的关键词列表
        """
        filtered_keywords = []
        for keyword in keywords:
            max_relevance = 0
            for theme in theme_words:
                relevance = self.compute_relevance(theme, keyword, context)
                max_relevance = max(max_relevance, relevance)

            if max_relevance >= self.threshold:
                filtered_keywords.append(keyword)

        return filtered_keywords
    
if __name__ == "__main__":
    # 上下文文本（例如平安银行年报内容）
    # context = (
    #     "平安银行的年度报告显示，财务增长稳健，净利润大幅提升。"
    #     "未来将继续加强对银行业务的数字化管理，扩大市场份额。"
    # )
    path = "/home/ubuntu/code/git/subject-word-extraction/scripts/clean_data/000001_2023_平安银行_2023年年度报告_2024-03-15.txt"
    with open(path,"r") as f:
        context = f.readlines()

    # 主题词和待筛选关键词
    # theme_words = ["财务", "增长", "银行"]  # 主题词
    theme_words = ["智能","智能制造"]  # 主题词

    # 初始化筛选器并设置相似度阈值
    filter = ChatGLM3KeywordFilter(threshold=0.5)
    keywords = filter.find_keywords(theme_words, tuple(context))
    print(keywords)


    # keywords = ["平安银行", "股东", "年度报告", "财务增长", "营业收入", "技术", "市场风险"]

    # # 初始化筛选器并设置相似度阈值
    # filter = ChatGLM3KeywordFilter(threshold=0.5)
    # filtered_keywords = filter.filter_keywords_by_theme_in_context(keywords, theme_words, context)

    # print("与主题词语义相似的关键词：", filtered_keywords)

